#include "session/manager.h"

#include "smart_serve_log.h"
#include "service/service.h"

namespace OHOS {
namespace SmartServe {

Session* SessionManager::GetOrCreateSession(EvalTask* task)
{
    switch (task->type_) {
        case EvalTask::Type::LLM: {
            return GetOrCreateSession(static_cast<EvalLLMTask*>(task));
        }
        default: {
            SMART_SERVE_LOGE("Unknown EvalTask type: %d", int(task->type_));
            return nullptr;
        }
    }
}

bool CheckAdapter(SessionAdapter* adapter)
{
    if (adapter->InitialResConfigs().empty()) {
        SMART_SERVE_LOGE("[SessionManager] Got empty InitialResConfigs for session %s", adapter->Name().c_str());
        return false;
    }

    return true;
}

void CleanUp(const std::vector<Model*>& models)
{
    auto service = Service::Instance();

    for (auto& model : models) {
        auto m = service->modelManager_.RemoveModel(model);
        m->AsyncTerminate();
        m->WaitForTerminated();
    }
}

Session* SessionManager::GetOrCreateSession(EvalLLMTask* task)
{
    const std::string& sessionName = task->modelName_;

    // 加管理锁
    auto service = Service::Instance();
    std::unique_lock<std::mutex> lock(service->msMgmtMutex_);

    // 通过 session 复用机制确认是否创建新 session
    auto existedSession = reuser_->GetSession(task);

    // 复用，不创建
    if (existedSession != nullptr) {
        return existedSession;
    }

    // 创建新的 Session

    // 拿到 SessionCreator，获取依赖的 model
    auto sc = GetCreator(sessionName);
    if (sc == nullptr) {
        return nullptr;
    }

    // 创建依赖 Model
    std::vector<Model*> depModels;
    std::vector<ModelAdapter*> depModelAdapters;
    for (auto& modelName : sc->DepModels()) {
        auto model = service->modelManager_.GetOrCreateModel(modelName);
        if (!model) {
            CleanUp(depModels);
            return nullptr;
        }

        depModels.push_back(model);
        depModelAdapters.push_back(model->Adapter());
    }

    // 创建 Session
    auto sa = sc->Create(depModelAdapters);
    if (sa == nullptr) {
        CleanUp(depModels);
        return nullptr;
    }
    if (!CheckAdapter(sa.get())) {
        sa->Finalize();
        CleanUp(depModels);
        return nullptr;
    }
    auto session = std::make_unique<Session>(std::move(sa));
    auto pSession = session.get();

    allSessions_.insert({sessionName, std::move(session)});

    // 链接 Model - Session
    for (auto& model : depModels) {
        pSession->AddModel(model);
        model->AddSession(pSession);
    }

    return pSession;
}

std::unique_ptr<Session> SessionManager::RemoveSession(Session* session)
{
    std::unique_ptr<Session> ret = nullptr;

    auto range = allSessions_.equal_range(session->Name());
    for (auto it = range.first; it != range.second; ++it) {
        if (it->second.get() == session) {
            ret = std::move(it->second);
            allSessions_.erase(it);
            break;
        }
    }

    return ret;
}

bool SessionManager::Register(std::unique_ptr<SessionCreator> creator)
{
    const auto& depModels = creator->DepModels();

    if (depModels.size() > 2) {
        for (std::size_t i = 1; i < depModels.size(); ++i) {
            if (depModels[i] < depModels[i - 1]) {
                SMART_SERVE_LOGE("[Register] depModels must be sorted in ascending lexicographical order by their names, "
                     "place %s before %s",
                     depModels[i].c_str(),
                     depModels[i - 1].c_str());
                return false;
            }
            if (depModels[i] == depModels[i - 1]) {
                SMART_SERVE_LOGE("[Register] depModels must be unique, found duplicates: %s", depModels[i].c_str());
                return false;
            }
        }
    }

    return registry_.insert({creator->Name(), std::move(creator)}).second;
}

} // namespace SmartServe
} // namespace OHOS